package cipher

import (
	"crypto/aes"
	"crypto/cipher"
)

type AesType int

const (
	AES_128_CBC AesType = 100 + iota
	AES_192_CBC
	AES_256_CBC
	AES_128_CFB
	AES_192_CFB
	AES_256_CFB
)

type AesCipher struct {
	Src []byte
	Key []byte
	T   AesType
	Iv  []byte //can be nil. default key[:blockSize]
}

// type to string
func (t AesType) String() string {
	switch t {
	case AES_128_CFB:
		return "AES_128_CFB"
	case AES_192_CFB:
		return "AES_192_CFB"
	case AES_256_CFB:
		return "AES_256_CFB"
	case AES_128_CBC:
		return "AES_128_CBC"
	case AES_192_CBC:
		return "AES_192_CBC"
	case AES_256_CBC:
		return "AES_256_CBC"
	default:
		panic("unknown aes cipher")
	}
}

// aes trim encrypt key
func (t AesType) getKey(key []byte) []byte {
	if len(key) < aes.BlockSize {
		panic("key too short")
	}
	switch t {
	case AES_128_CFB, AES_128_CBC:
		return key[:aes.BlockSize]
	case AES_192_CFB, AES_192_CBC:
		return key[:24]
	case AES_256_CFB, AES_256_CBC:
		return key[:32]
	default:
		panic("unknown cipher: " + t.String())
	}
}

// aes iv
func (c *AesCipher) getIv() []byte {
	if c.Iv == nil || len(c.Iv) < c.blockSize() {
		return c.Key[:c.blockSize()]
	}
	return c.Iv[:c.blockSize()]
}

// aes block size is 16
func (c *AesCipher) blockSize() int {
	return aes.BlockSize
}

// aes encrypt
func (c *AesCipher) Encrypt() ([]byte, error) {
	key := c.T.getKey(c.Key)
	// new block
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, CError{
			detail: err.Error(),
		}
	}
	iv := c.Key[:block.BlockSize()]
	var dst []byte

	switch c.T {
	case AES_128_CFB, AES_192_CFB, AES_256_CFB:
		dst = make([]byte, len(c.Src))
		stream := cipher.NewCFBEncrypter(block, iv)
		stream.XORKeyStream(dst, c.Src)
	case AES_128_CBC, AES_192_CBC, AES_256_CBC:
		c.Src = PKCS7Padding(c.Src, block.BlockSize())
		dst = make([]byte, len(c.Src))
		blockMode := cipher.NewCBCEncrypter(block, iv)
		blockMode.CryptBlocks(dst, c.Src)
	default:
		return nil, CError{
			detail: "can't encrypt with " + c.T.String(),
		}
	}
	return dst, nil
}

// aes decrypt
func (c *AesCipher) Decrypt() ([]byte, error) {
	key := c.T.getKey(c.Key)
	// new block
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, CError{
			detail: err.Error(),
		}
	}
	iv := c.Key[:block.BlockSize()]
	dst := make([]byte, len(c.Src))

	switch c.T {
	case AES_128_CFB, AES_192_CFB, AES_256_CFB:
		stream := cipher.NewCFBDecrypter(block, iv)
		stream.XORKeyStream(dst, c.Src)
	case AES_128_CBC, AES_192_CBC, AES_256_CBC:
		blockMode := cipher.NewCBCDecrypter(block, iv)
		blockMode.CryptBlocks(dst, c.Src)
		dst = PKCS7UnPadding(dst)
	default:
		return nil, CError{
			detail: "can't decrypt with " + c.T.String(),
		}
	}
	return dst, nil
}
